# -*- coding: utf-8 -*-
import argparse, json, numpy as np, cv2
# from rknn.api import RKNN
from rknnlite.api import RKNNLite

def imread_gray_any(p):
    data = np.fromfile(p, dtype=np.uint8)
    img = cv2.imdecode(data, 0)
    if img is None: raise FileNotFoundError(p)
    return img

def imread_gray(img):
    return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


def to_u8_from_any(y):
    """
    尝试把 RKNN 输出(可能是 float32/int8/uint8)稳定转换为 0..255 的 uint8 图
    约定：训练/ONNX阶段输出范围在[-1,1]（我们就按这个路径优先判断）
    """
    arr = y.squeeze()
    if arr.dtype == np.float32 or arr.dtype == np.float64:
        vmin, vmax = float(arr.min()), float(arr.max())
        # 大多数情况下是[-1,1]
        if vmin >= -1.5 and vmax <= 1.5:
            out = ((arr + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
        else:
            # fallback：线性拉伸到0..255
            m, M = arr.min(), arr.max()
            out = ((arr - m) / (M - m + 1e-12) * 255.0).clip(0, 255).astype(np.uint8)
    elif arr.dtype == np.int8:
        # 粗略视作对称量化，[-128,127] -> [-1,1] -> 0..255
        out = (((arr.astype(np.float32) / 127.0) + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
    elif arr.dtype == np.uint8:
        # 已经是图像量级（多数在0..255）
        out = arr.astype(np.uint8)
    else:
        # 其它情况统一归一化
        m, M = arr.min(), arr.max()
        out = ((arr - m) / (M - m + 1e-12) * 255.0).clip(0, 255).astype(np.uint8)
    return out

class remix_model:
    def __init__(self):
        self.rknn=RKNNLite()

        


    def remix(self,vis,ir):
        Rknn='/root/remix_cam/rknn_test/fusion_model.rknn'
        device ='cpu'
        config='/root/remix_cam/rknn_test/config.json'
        H, W = json.load(open(config,'r'))['input_size']
        vis = cv2.resize(imread_gray(vis), (W,H), interpolation=cv2.INTER_AREA)
        ir  = cv2.resize(imread_gray(ir),  (W,H), interpolation=cv2.INTER_AREA)
        x = np.stack([ir, vis], axis=0).astype(np.float32) / 127.5 - 1.0
        x = x[None, ...]  # (1,2,H,W)
        rknn = self.rknn
        assert rknn.load_rknn(Rknn) == 0, 'load_rknn failed'
        if device == 'cpu':
            assert rknn.init_runtime() == 0, 'init_runtime failed'
        else:
            assert rknn.init_runtime(target=device) == 0, 'init_runtime failed'

        outs = rknn.inference(inputs=[x])
        y = outs[0]  # 形状一般是 (1,1,H,W) 或类似
        print('[INFO] RKNN output:', y.dtype, y.shape, 'min/max=', float(y.min()), float(y.max()))

        y8 = to_u8_from_any(y)
        cv2.imencode('.png', y8)[1].tofile('./test.png')
        return y8

    def release(self):
        """释放资源"""
        try:
            self.rknn.release()
        except Exception as e:
            print(f"释放RKNN失败: {e}")
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.release()
    
    def __del__(self):
        """析构函数"""
        self.release()
        


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--rknn', default='/root/remix_cam/rknn_test/fusion_model.rknn')
    ap.add_argument('--vis', default='/root/vis.jpg')
    ap.add_argument('--ir', default='/root/ir.jpg')
    ap.add_argument('--out', default='./fused_rknn.png')
    ap.add_argument('--config', default='/root/remix_cam/rknn_test/config.json')
    ap.add_argument('--device', default='cpu', choices=['cpu','rk3588','rk3568','rk3399pro'])
    args = ap.parse_args()

    H, W = json.load(open(args.config,'r'))['input_size']

    # 读图并预处理到[-1,1]，堆成(1,2,H,W)
    vis = cv2.resize(imread_gray_any(args.vis), (W,H), interpolation=cv2.INTER_AREA)
    ir  = cv2.resize(imread_gray_any(args.ir),  (W,H), interpolation=cv2.INTER_AREA)
    x = np.stack([ir, vis], axis=0).astype(np.float32) / 127.5 - 1.0
    x = x[None, ...]  # (1,2,H,W)

    rknn = RKNNLite()
    assert rknn.load_rknn(args.rknn) == 0, 'load_rknn failed'
    if args.device == 'cpu':
        assert rknn.init_runtime() == 0, 'init_runtime failed'
    else:
        assert rknn.init_runtime(target=args.device) == 0, 'init_runtime failed'

    outs = rknn.inference(inputs=[x])
    y = outs[0]  # 形状一般是 (1,1,H,W) 或类似
    print('[INFO] RKNN output:', y.dtype, y.shape, 'min/max=', float(y.min()), float(y.max()))

    y8 = to_u8_from_any(y)
    cv2.imencode('.png', y8)[1].tofile(args.out)
    print(f'[OK] saved -> {args.out}')
    rknn.release()

if __name__ == '__main__':
    main()
